#include "opencv2/opencv.hpp"
#include "iostream"
#include "cmath"

using namespace std;

int min_vec_val(cv::Vec3b &rgb) {
    int min_val = rgb[0];
    min_val = (min_val > rgb[1]) ? rgb[1] : min_val;
    min_val = (min_val > rgb[2]) ? rgb[2] : min_val;
    return min_val;
}

double min_vec_val(double a, double b, double c) {
    double min_val = a;
    min_val = (min_val > b) ? b : min_val;
    min_val = (min_val > c) ? c : min_val;
    return min_val;
}

void dehaze(cv::Mat &in, cv::Mat &out, int size, double w = 0.95, double t0 = 0.1) {
    cv::Mat kernel = cv::getStructuringElement(cv::MORPH_RECT, cv::Size(size, size));
    int height = in.size().height;
    int width = in.size().width;
    int search_pts = ceil(height * width * 0.001);
    
    // cal dark channel
    cv::Mat dark_image = cv::Mat::zeros(height, width, CV_8UC1);
    for (int i = 0; i < height; ++ i) {
        for (int j = 0; j < width; ++ j) {
            dark_image.at<uint8_t>(i, j) = min_vec_val(in.at<cv::Vec3b>(i, j));
        }
    }
    cv::erode(dark_image, dark_image, kernel);

    // estimate A
    int pix_counts[256];
    uint8_t ktop_index = 255;
    memset(pix_counts, 0, sizeof(pix_counts));
    for (int i = 0; i < height; ++ i) {
        for (int j = 0; j < width; ++ j) {
            pix_counts[dark_image.at<uint8_t>(i, j)] ++;
        }
    }
    
    while (search_pts > 0) {
        search_pts -= pix_counts[ktop_index --];
    }
    // then we choose the pixel whose value > ktop_index and update the final A
    int A[3];
    cv::Vec3i cur_bgr;
    A[0] = A[1] = A[2] = 1;
    for (int i = 0; i < height; ++ i) {
        for (int j = 0; j < width; ++ j) {
            if (dark_image.at<uint8_t>(i, j) > ktop_index) {
                cur_bgr = in.at<cv::Vec3b>(i, j);
                A[0] = max(A[0], cur_bgr[0]);
                A[1] = max(A[1], cur_bgr[1]);
                A[2] = max(A[2], cur_bgr[2]);
            }
        }
    }

    // estimate t
    cv::Mat imageA = cv::Mat::zeros(height, width, CV_64FC1), t;
    double cur_min;
    for (int i = 0; i < height; ++ i) {
        for (int j = 0; j < width; ++ j) {
            cur_bgr = in.at<cv::Vec3b>(i, j);
            cur_min = min_vec_val(1.0 * cur_bgr[0] / A[0], 1.0 * cur_bgr[1] / A[1], 1.0 * cur_bgr[2] / A[2]);
            imageA.at<double>(i, j) = cur_min;       // don't forget to normalise
        }
    }
    cv::erode(imageA, imageA, kernel);
    imageA.convertTo(t, CV_64FC1, -1. * w, 1);

    double cur_t;
    cv::Vec3d temp_pixel, cur_pixel;
    cv::Mat result = cv::Mat::zeros(height, width, CV_8UC3);

    for (int i = 0; i < height; ++ i) {
        for (int j = 0; j < width; ++ j) {
            cur_t = max(t.at<double>(i, j), t0);
            cur_pixel = in.at<cv::Vec3b>(i, j);
            result.at<cv::Vec3b>(i, j)[0] = ((cur_pixel[0] - A[0]) / cur_t + A[0]);
            result.at<cv::Vec3b>(i, j)[1] = ((cur_pixel[1] - A[1]) / cur_t + A[1]);
            result.at<cv::Vec3b>(i, j)[2] = ((cur_pixel[2] - A[2]) / cur_t + A[2]);
        }
    }

    result.convertTo(out, CV_8UC3);
}

int main(int argc, char const *argv[])
{
    if (argc != 6) {
        cout << "useage: "<<  argv[0] << " <input_image> <output_image> size w t0" << endl;
        exit(-1);
    }
    
    cv::Mat in, out;
    in = cv::imread(argv[1]);
    dehaze(in, out, atoi(argv[3]), atof(argv[4]), atof(argv[5]));
    cv::imwrite(argv[2], out);

    return 0;
}
